# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# Adapted from
# https://github.com/huggingface/transformers/blob/main/src/transformers/utils/versions.py#L49
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib.metadata
import operator
import re
import sys
from typing import Optional
from packaging import version


# Operator mapping for version comparison
_ops_map = {
    "<": operator.lt,
    "<=": operator.le,
    "==": operator.eq,
    "!=": operator.ne,
    ">=": operator.ge,
    ">": operator.gt,
}


def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint):
    if op not in _ops_map:
        raise ValueError(f"{requirement}: need one of {list(_ops_map.keys())}, but got {op}")
    if got_ver is None or want_ver is None:
        raise ValueError(
            f"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider"
            f" reinstalling {pkg}."
        )
    if not _ops_map[op](version.parse(got_ver), version.parse(want_ver)):
        raise ImportError(
            f"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}"
        )


def require_version(requirement: str, hint: Optional[str] = None) -> None:
    """
    Perform a runtime check of the dependency versions, using the exact same syntax used by pip.

    The installed module version comes from the *site-packages* dir via *importlib.metadata*.

    Args:
        requirement (`str`): pip style definition, e.g.,  "tokenizers==0.9.4", "tqdm>=4.27", "numpy"
        hint (`str`, *optional*): what suggestion to print in case of requirements not being met

    Example:

    ```python
    require_version("pandas>1.1.2")
    require_version("numpy>1.18.5", "this is important to have for whatever reason")
    ```
    """

    hint = f"\n{hint}" if hint is not None else ""

    # non-versioned check
    if re.match(r"^[\w_\-\d]+$", requirement):
        pkg, op, want_ver = requirement, None, None
    else:
        match = re.findall(r"^([^!=<>\s]+)([\s!=<>]{1,2}.+)", requirement)
        if not match:
            raise ValueError(
                "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"
                f" got {requirement}"
            )
        pkg, want_full = match[0]
        want_range = want_full.split(",")  # there could be multiple requirements
        wanted = {}
        for w in want_range:
            match = re.findall(r"^([\s!=<>]{1,2})(.+)", w)
            if not match:
                raise ValueError(
                    "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"
                    f" but got {requirement}"
                )
            op, want_ver = match[0]
            wanted[op] = want_ver

    # special case
    if pkg == "python":
        got_ver = ".".join([str(x) for x in sys.version_info[:3]])
        for op, want_ver in wanted.items():
            _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)
        return

    # check if any version is installed
    try:
        got_ver = importlib.metadata.version(pkg)
    except importlib.metadata.PackageNotFoundError as e:
        raise importlib.metadata.PackageNotFoundError(
            f"The '{requirement}' distribution was not found and is required by this application. {hint}"
        ) from e

    # check that the right version is installed if version number or a range was provided
    if want_ver is not None:
        for op, want_ver in wanted.items():
            _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)


def _parse_requirement_string(requirement_string: str):
    """
    Parses a requirement string into a package name and a list of conditions.

    Args:
        requirement_string (str): The requirement string, e.g.,
                                  "transformers>=4.45.0,<4.50.0" or "numpy" or "python==3.9.1".

    Returns:
        tuple: (pkg_name, list_of_conditions) or (None, None) on parse failure.
               list_of_conditions is a list of (operator_str, version_str) tuples.
               Returns (pkg_name, []) if no version is specified (existence check).
    """
    cleaned_requirement = requirement_string.strip()
    if not cleaned_requirement:
        return None, None  # Empty requirement string

    # Case 1: Only package name (no version specifiers)
    # Matches common package name characters: alphanumeric, underscore, dot, hyphen.
    simple_pkg_match = re.fullmatch(r"^[a-zA-Z0-9_.\-]+$", cleaned_requirement)
    if simple_pkg_match:
        return cleaned_requirement, []

    # Case 2: Package name with version specifiers
    # Regex to separate package name from the rest of the version specifiers.
    # Group 1: package_name (e.g., "transformers")
    # Group 2: the entire version specification string (e.g., ">=4.45.0,<4.50.0" or " ==3.9.1 ")
    # Package names can contain letters, numbers, '.', '_', '-'.
    # Version specifiers must start with an operator.
    match_pkg_and_spec = re.match(r"^\s*([a-zA-Z0-9_.\-]+)\s*([!=<>]{1,2}.*)\s*$", cleaned_requirement)
    if not match_pkg_and_spec:
        return None, None

    pkg_name = match_pkg_and_spec.group(1)
    version_specs_str = match_pkg_and_spec.group(2).strip()

    if not pkg_name:  # Should not happen with the regex structure
        return None, None

    conditions = []
    # Split by comma for multiple conditions (e.g., ">=1.0, <2.0")
    spec_parts = [s.strip() for s in version_specs_str.split(",") if s.strip()]

    if not spec_parts and version_specs_str:  # e.g. "pkg," or "pkg==,"
        return None, None

    for part in spec_parts:
        # Match operator and version for each part
        # Group 1: operator (e.g., ">=", "<")
        # Group 2: version_string (e.g., "4.45.0")
        # Allows spaces around operator and version, which are then stripped.
        # Version string is matched non-greedily to handle various characters.
        match_condition = re.match(r"^\s*([!=<>]{1,2})\s*(.+?)\s*$", part)
        if not match_condition:
            return None, None  # Malformed condition part

        op_str = match_condition.group(1).strip()
        ver_str = match_condition.group(2).strip()

        if op_str not in _ops_map:
            return None, None  # Invalid operator
        if not ver_str:
            return None, None  # Empty version string for an operator

        conditions.append((op_str, ver_str))

    # If version_specs_str was present but no conditions were parsed (e.g. "pkg >=")
    if not conditions and version_specs_str:
        return None, None

    return pkg_name, conditions


def check_package_version(requirement_string: str) -> bool:
    """
    Checks if an installed package meets the specified version requirements.

    Args:
        requirement_string (str): A requirement string in a pip-like format,
                                  e.g., "transformers<=4.51.3, >=4.48.0, !=4.51.2", "numpy==1.20.0",
                                  "requests", "python>=3.8".

    Returns:
        bool: True if the installed version meets all requirements, False otherwise
              (including if package is not found, requirement is malformed,
              or version comparison fails).
    """
    pkg_name, conditions = _parse_requirement_string(requirement_string)

    if pkg_name is None:
        # Malformed requirement string or unhandled parsing case.
        return False

    installed_version_str: str
    # Special case for Python version
    if pkg_name.lower() == "python":
        installed_version_str = ".".join(map(str, sys.version_info[:3]))
    else:
        # Get installed package version
        try:
            installed_version_str = importlib.metadata.version(pkg_name)
        except importlib.metadata.PackageNotFoundError:
            # Package not found. If there were any conditions or even just an
            # existence check (empty conditions list), this is a failure.
            return False

    # If there are no conditions, it's an existence check.
    # If we reached here, the package (or Python) exists.
    if not conditions:
        return True

    # Parse the installed version string
    try:
        parsed_installed_version = version.parse(installed_version_str)
    except version.InvalidVersion:
        # Installed version string is invalid, cannot reliably compare.
        return False

    # Check all conditions
    for op_str, required_version_str in conditions:
        try:
            parsed_required_version = version.parse(required_version_str)
        except version.InvalidVersion:
            # Requirement's version string is invalid.
            return False  # Requirement itself is flawed

        comparison_func = _ops_map[op_str]

        # Perform the comparison
        if not comparison_func(parsed_installed_version, parsed_required_version):
            return False  # This specific condition is not met

    # All conditions were met
    return True


def is_transformers_version_equal_to_4_46():
    return check_package_version("transformers>=4.46.0, <= 4.46.1")
